import os
import sys
import subprocess
import random
import re
import requests
from pathlib import Path
from http import HTTPStatus
from dashscope import VideoSynthesis
import dashscope
from dotenv import load_dotenv
import time

# 加载.env文件中的环境变量
load_dotenv()

# 设置API密钥和基础URL
api_key = os.getenv("DASHSCOPE_API_KEY")
dashscope.base_http_api_url = 'https://dashscope.aliyuncs.com/api/v1'

# 检查API密钥是否存在
if not api_key:
    raise ValueError("未找到DASHSCOPE_API_KEY环境变量，请检查.env文件配置")

# 全局变量用于存储任务引用
current_task_ref = None


def set_current_task(task_ref):
    """
    设置当前任务引用，用于实时更新日志
    """
    global current_task_ref
    current_task_ref = task_ref


def update_task_log(log_message):
    """
    实时更新任务日志
    """
    global current_task_ref
    if current_task_ref and 'log' in current_task_ref:
        # 添加新日志行
        if current_task_ref['log']:
            current_task_ref['log'] += '\n' + log_message
        else:
            current_task_ref['log'] = log_message


def get_audio_duration(audio_path):
    """
    获取音频/视频文件的时长(秒)
    """
    # 确保路径是字符串格式
    audio_path_str = str(audio_path)
    
    cmd = [
        'ffprobe', '-v', 'error', '-show_entries', 'format=duration',
        '-of', 'default=noprint_wrappers=1:nokey=1', audio_path_str
    ]
    
    try:
        result = subprocess.run(
            cmd, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE, 
            text=True, 
            timeout=30,
            encoding='utf-8',
            errors='replace'  # 替换无法解码的字符
        )
        if result.returncode != 0:
            print(f"ffprobe执行失败 {audio_path}: {result.stderr}")
            return None
            
        duration_str = result.stdout.strip()
        if not duration_str:
            print(f"未能获取到时长信息 {audio_path}")
            return None
            
        duration = float(duration_str)
        return min(duration, 30.0)  # 限制最大时长为30秒，给更多宽容度
    except ValueError as e:
        print(f"时长转换失败 {audio_path}: {e}, 原始输出: '{result.stdout.strip()}'")
        return None
    except Exception as e:
        print(f"获取媒体时长失败 {audio_path}: {e}")
        return None


def extract_audio_from_video(video_path, audio_path, duration):
    """
    从视频中提取指定时长的音频
    """
    cmd = [
        'ffmpeg', '-y', '-i', str(video_path),
        '-t', str(duration), '-vn', '-acodec', 'copy', str(audio_path)
    ]
    
    # 添加重试机制，最多尝试3次
    max_retries = 3
    for attempt in range(max_retries):
        try:
            result = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=60,
                encoding='utf-8',
                errors='replace'
            )
            if result.returncode == 0:
                return True
            else:
                print(f"提取音频失败 (尝试 {attempt + 1}/{max_retries}): {result.stderr}")
        except Exception as e:
            print(f"提取音频出错 (尝试 {attempt + 1}/{max_retries}): {e}")
        
        # 如果不是最后一次尝试，等待一段时间再重试
        if attempt < max_retries - 1:
            time.sleep(2)
    
    return False


def merge_audio_video(video_path, audio_path, output_path):
    """
    将音频和视频合并
    """
    cmd = [
        'ffmpeg', '-y', '-i', str(video_path),
        '-i', str(audio_path), '-c:v', 'copy', '-c:a', 'aac', str(output_path)
    ]
    
    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=60,
            encoding='utf-8',
            errors='replace'
        )
        return result.returncode == 0
    except Exception as e:
        print(f"合并音视频失败: {e}")
        return False


def transcode_video_for_recording(input_path, output_path):
    """
    将视频转码为更适合录制的格式
    
    Args:
        input_path (str): 输入视频路径
        output_path (str): 输出视频路径
        
    Returns:
        bool: 转码是否成功
    """
    try:
        # 检查输入文件是否存在
        if not os.path.exists(input_path):
            print(f"输入文件不存在: {input_path}")
            return False
            
        # 构造FFmpeg命令
        cmd = [
            'ffmpeg',
            '-y',  # 覆盖输出文件
            '-i', str(input_path),  # 输入文件
            '-vf', 'setpts=PTS+0.05/TB,eq=brightness=0.02:contrast=1.03,boxblur=1:1',  # 视频滤镜
            '-af', 'atempo=1.0',  # 音频滤镜
            '-c:v', 'libx264',  # 视频编码器
            '-preset', 'medium',  # 编码预设
            '-crf', '23',  # 视频质量
            '-c:a', 'aac',  # 音频编码器
            '-b:a', '128k',  # 音频比特率
            '-movflags', '+faststart',  # 快速开始标志
            '-map_metadata', '-1',  # 移除元数据
            str(output_path)
        ]
        
        # 执行转码
        print(f"正在执行视频转码: {' '.join([str(c) for c in cmd])}")
        result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8')
        
        if result.returncode == 0:
            print(f"视频转码成功: {output_path}")
            return True
        else:
            print(f"视频转码失败: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"视频转码出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def async_call_i2v(img_url, prompt, resolution="1080P", negative_prompt=None):
    """
    异步调用图片到视频生成功能
    """
    # 异步调用，返回一个task_id
    # 准备调用参数
    call_params = {
        'api_key': api_key,
        'model': 'wan2.2-i2v-flash',
        'prompt': prompt,
        'resolution': resolution,
        'img_url': img_url
    }
    
    # 如果提供了negative_prompt，则添加到参数中
    if negative_prompt is not None:
        call_params['negative_prompt'] = negative_prompt
    
    rsp = VideoSynthesis.async_call(**call_params)
    
    if rsp.status_code == HTTPStatus.OK:
        return rsp.output.task_id, None
    else:
        error_msg = f'Failed, status_code: {rsp.status_code}, code: {rsp.code}, message: {rsp.message}'
        return None, error_msg


def fetch_task_status(task_id):
    """
    获取异步任务状态
    """
    # 直接使用任务ID获取任务状态，提供API密钥
    status = VideoSynthesis.fetch(task_id, api_key=api_key)
    
    if status.status_code == HTTPStatus.OK:
        return status.output.task_status, status.output.video_url if hasattr(status.output, 'video_url') else None, None
    else:
        error_msg = f'Failed, status_code: {status.status_code}, code: {status.code}, message: {rsp.message}'
        return None, None, error_msg


def wait_for_task(task_id):
    """
    等待异步任务完成
    """
    # 直接使用任务ID等待任务完成，提供API密钥
    rsp = VideoSynthesis.wait(task_id, api_key=api_key)
    
    if rsp.status_code == HTTPStatus.OK:
        return rsp.output.video_url, None
    else:
        error_msg = f'Failed, status_code: {rsp.status_code}, code: {rsp.code}, message: {rsp.message}'
        return None, error_msg


def process_single_image_with_audio(image_path, prompt, resolution, audio_files, output_path, negative_prompt=None, duration=None):
    """
    处理单个图片，生成视频并添加音频
    
    Args:
        image_path (str): 图片路径
        prompt (str): 视频生成提示
        resolution (str): 分辨率
        audio_files (list): 音频文件列表
        output_path (str): 输出视频路径
        negative_prompt (str, optional): 负面提示
        duration (float, optional): 固定时长（秒），默认为None表示自动检测
    """
    logs = []
    
    # 构造file://协议的路径
    if os.name == 'nt':  # Windows系统
        # 将反斜杠替换为正斜杠
        image_path = image_path.replace('\\', '/')
    img_url = "file://" + image_path
    
    logs.append(f"开始处理图片: {image_path}")
    logs.append(f"视频描述: {prompt}")
    logs.append(f"分辨率: {resolution}")
    if negative_prompt:
        logs.append(f"负面描述: {negative_prompt}")
    
    # 异步调用图片到视频生成功能
    task_id, error = async_call_i2v(img_url, prompt, resolution, negative_prompt)
    if error:
        logs.append(f"调用失败: {error}")
        return False, "\n".join(logs)
    
    logs.append(f"任务已提交，任务ID: {task_id}")
    
    # 等待任务完成
    video_url, error = wait_for_task(task_id)
    if error:
        logs.append(f"任务执行失败: {error}")
        return False, "\n".join(logs)
    
    logs.append(f"视频生成完成: {video_url}")
    
    # 下载视频
    import requests
    video_data = requests.get(video_url)
    if video_data.status_code != 200:
        logs.append(f"下载视频失败，状态码: {video_data.status_code}")
        return False, "\n".join(logs)
    
    # 保存视频到临时文件
    temp_video = Path(output_path).with_suffix('.temp.mp4')
    try:
        with open(temp_video, 'wb') as f:
            f.write(video_data.content)
        logs.append(f"视频已下载到: {temp_video}")
    except Exception as e:
        logs.append(f"保存视频失败: {e}")
        return False, "\n".join(logs)
    
    # 检查临时文件是否存在且不为空
    if not temp_video.exists() or temp_video.stat().st_size == 0:
        logs.append("视频文件下载失败或为空")
        # 清理可能存在的空文件
        if temp_video.exists():
            try:
                temp_video.unlink()
            except:
                pass
        return False, "\n".join(logs)
    
    # 处理音频文件
    if audio_files:
        # 随机选择一个音频文件
        selected_audio = random.choice(audio_files)
        logs.append(f"随机选择音频文件: {selected_audio}")
        
        # 获取视频时长
        if duration is not None:
            # 使用提供的固定时长
            video_duration = duration
            logs.append(f"使用指定视频时长: {video_duration}秒")
        else:
            # 自动检测视频时长
            video_duration = get_audio_duration(str(temp_video))
            if video_duration:
                logs.append(f"检测到视频时长: {video_duration}秒")
            else:
                logs.append("无法获取视频时长")
    
    if audio_files:
        # 随机选择一个音频文件
        selected_audio = random.choice(audio_files)
        logs.append(f"随机选择音频文件: {selected_audio}")
        
        # 提取与视频时长相同的音频片段，添加重试机制
        temp_audio = Path(output_path).with_suffix('.temp_audio.mp3')
        
        # 尝试提取音频，如果失败则尝试其他音频文件
        audio_extracted = False
        tried_audios = []  # 记录已经尝试过的音频文件
        
        # 先尝试选中的音频文件
        if extract_audio_from_video(selected_audio, str(temp_audio), video_duration):
            logs.append(f"音频片段提取完成: {temp_audio}")
            audio_extracted = True
        else:
            logs.append(f"从选中音频文件提取失败: {selected_audio}")
        
        tried_audios.append(selected_audio)
        
        # 如果第一次尝试失败，尝试其他音频文件
        if not audio_extracted:
            # 打乱音频文件列表，尝试其他音频
            other_audios = [audio for audio in audio_files if audio not in tried_audios]
            random.shuffle(other_audios)
            
            # 最多尝试3个不同的音频文件
            max_audio_retries = min(3, len(other_audios))
            for i in range(max_audio_retries):
                alternative_audio = other_audios[i]
                logs.append(f"尝试其他音频文件 ({i+1}/{max_audio_retries}): {alternative_audio}")
                
                if extract_audio_from_video(alternative_audio, str(temp_audio), video_duration):
                    logs.append(f"音频片段提取完成: {temp_audio}")
                    audio_extracted = True
                    break
                else:
                    logs.append(f"从替代音频文件提取失败: {alternative_audio}")
                    tried_audios.append(alternative_audio)
        
        if audio_extracted:
            # 合并音视频
            if merge_audio_video(str(temp_video), str(temp_audio), output_path):
                logs.append(f"音视频合并完成: {output_path}")
                result = True
            else:
                logs.append("音视频合并失败")
                result = False
        else:
            logs.append("所有音频文件尝试提取都失败了")
            result = False
            
        # 清理临时音频文件
        if temp_audio.exists():
            try:
                temp_audio.unlink()
            except:
                pass
            
        # 清理临时视频文件（如果存在且处理成功）
        if result and temp_video.exists():
            try:
                temp_video.unlink()
            except:
                pass
            
        return result, "\n".join(logs)
    else:
        # 没有音频文件，直接将生成的视频作为最终输出
        try:
            temp_video.rename(output_path)
            logs.append(f"视频已保存到: {output_path}")
            return True, "\n".join(logs)
        except Exception as e:
            # 如果rename失败，尝试使用shutil.move
            try:
                import shutil
                shutil.move(str(temp_video), output_path)
                logs.append(f"视频已保存到: {output_path}")
                return True, "\n".join(logs)
            except Exception as e2:
                logs.append(f"保存视频失败: {e} (重试也失败: {e2})")
                # 清理临时文件
                if temp_video.exists():
                    try:
                        temp_video.unlink()
                    except:
                        pass
                return False, "\n".join(logs)

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python script.py <image_folder> <output_folder> <prompt>")
        sys.exit(1)

    image_folder = Path(sys.argv[1])
    output_folder = Path(sys.argv[2])
    prompt = sys.argv[3]

    if not image_folder.exists() or not image_folder.is_dir():
        print(f"Image folder '{image_folder}' does not exist or is not a directory.")
        sys.exit(1)

    if not output_folder.exists():
        output_folder.mkdir(parents=True)

    image_files = list(image_folder.glob('*.*'))
    audio_files = list(Path('audio').glob('*.*'))

    if not image_files:
        print("No image files found in the image folder.")
        sys.exit(1)

    for image_file in image_files:
        output_file = output_folder / image_file.name
        output_file = output_file.with_suffix('.mp4')
        print(f"Processing {image_file} -> {output_file}")
        # 处理单个图片
        result, process_logs = process_single_image_with_audio(
            str(image_file), prompt, "1080P", audio_files, str(output_file), duration=5.0
        )
        if result:
            print(f"Processed {image_file} -> {output_file}")
        else:
            print(f"Failed to process {image_file}: {process_logs}")
